// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/gtest_utils.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/sparksql/aggregates/register.h>

using namespace kumo::pollux::functions::aggregate::test;

namespace kumo::pollux::functions::aggregate::sparksql::test {

namespace {

class MinMaxByAggregateTest : public AggregationTestBase {
 protected:
  void SetUp() override {
    AggregationTestBase::SetUp();
    registerAggregateFunctions("spark_");
  }
};

TEST_F(MinMaxByAggregateTest, maxBy) {
  auto vectors = {make_row_vector({
      make_flat_vector<int32_t>({1, 2, 3}),
      make_flat_vector<int32_t>({11, 12, 12}),
  })};

  auto expected = {make_row_vector({
      make_flat_vector<int32_t>(std::vector<int32_t>({3})),
  })};

  testAggregations(vectors, {}, {"spark_max_by(c0, c1)"}, expected);
}

TEST_F(MinMaxByAggregateTest, minBy) {
  auto vectors = {make_row_vector({
      make_flat_vector<int32_t>({1, 2, 3}),
      make_flat_vector<int32_t>({12, 11, 11}),
  })};

  auto expected = {make_row_vector({
      make_flat_vector<int32_t>(std::vector<int32_t>({3})),
  })};

  testAggregations(vectors, {}, {"spark_min_by(c0, c1)"}, expected);
}

TEST_F(MinMaxByAggregateTest, arrayCompare) {
  auto data = make_row_vector({
      make_array_vector<int64_t>({
          {1, 2, 3},
          {4, 5},
          {6, 7, 8},
      }),
      make_nullable_array_vector<int64_t>({
          {4, 5},
          {std::nullopt, 2},
          {6, 7, 8},
      }),
  });

  auto expected = make_row_vector({
      make_array_vector<int64_t>({
          {4, 5},
      }),
      make_array_vector<int64_t>({
          {6, 7, 8},
      }),
  });

  testAggregations(
      {data}, {}, {"spark_min_by(c0, c1)", "spark_max_by(c0, c1)"}, {expected});

  data = make_row_vector({
      make_array_vector<int64_t>({
          {1, 2, 3},
          {4, 5},
          {6, 7, 8},
      }),
      make_nullable_array_vector<int64_t>({
          {1, 2, 3},
          {3, std::nullopt, 4},
          {6, 7, 8},
      }),
  });

  expected = make_row_vector({
      make_array_vector<int64_t>({
          {1, 2, 3},
      }),
      make_array_vector<int64_t>({
          {6, 7, 8},
      }),
  });

  testAggregations(
      {data}, {}, {"spark_min_by(c0, c1)", "spark_max_by(c0, c1)"}, {expected});
}

TEST_F(MinMaxByAggregateTest, mapCompare) {
  auto data = make_row_vector({
      make_array_vector<int64_t>({}),
      make_nullable_map_vector<int64_t, int64_t>({}),
  });
  std::vector<RowVectorPtr> expected = {};

  POLLUX_ASSERT_USER_THROW(
      testAggregations({data}, {}, {"spark_min_by(c0, c1)"}, expected),
      "Aggregate function signature is not supported");

  POLLUX_ASSERT_USER_THROW(
      testAggregations({data}, {}, {"spark_max_by(c0, c1)"}, expected),
      "Aggregate function signature is not supported");
}

TEST_F(MinMaxByAggregateTest, rowCompare) {
  auto data = make_row_vector({
      make_array_vector<int64_t>({
          {1, 2, 3},
          {4, 5},
          {6, 7, 8},
      }),
      make_row_vector({make_nullable_flat_vector<int32_t>({
          1,
          std::nullopt,
          3,
      })}),
  });

  auto expected = make_row_vector({
      make_array_vector<int64_t>({
          {4, 5},
      }),
      make_array_vector<int64_t>({
          {6, 7, 8},
      }),
  });

  testAggregations(
      {data}, {}, {"spark_min_by(c0, c1)", "spark_max_by(c0, c1)"}, {expected});

  data = make_row_vector({
      make_array_vector<int64_t>({
          {1, 2, 3},
          {4, 5},
          {6, 7, 8},
      }),
      make_row_vector({make_nullable_flat_vector<int32_t>({
          1,
          2,
          3,
      })}),
  });

  expected = make_row_vector({
      make_array_vector<int64_t>({
          {1, 2, 3},
      }),
      make_array_vector<int64_t>({
          {6, 7, 8},
      }),
  });

  testAggregations(
      {data}, {}, {"spark_min_by(c0, c1)", "spark_max_by(c0, c1)"}, {expected});
}

} // namespace
} // namespace kumo::pollux::functions::aggregate::sparksql::test
